import os
import sys
import logging
import numpy as np

def Test_T5():
    model_dir = "./models/t5-small"
    if not os.path.exists(model_dir):
        raise ValueError(f"model dir:{model_dir} is not exist.")

    contents = "summarize: studies have shown that owning a dog is good for you "
    
    ## Torch infer
    from src.t5.torch_infer import TorchInfer
    sample = TorchInfer(model_dir=model_dir)
    torch_res = sample.infer(contents)

    ## ONNX infer
    from src.t5.onnx_infer import ONNXInfer
    sample = ONNXInfer(model_dir=model_dir)
    logging.debug(f"{ONNXInfer.__mro__}")
    onnx_res = sample.infer(contents)

    ## check result data
    if (onnx_res == torch_res).all():
        logging.info(f"t5 check result ok")

def Test_electra():
    model_dir = "./models/electra"
    if not os.path.exists(model_dir):
        raise ValueError(f"model dir:{model_dir} is not exist.")

    contents = "Hello, my dog is cute"
    
    ## Torch infer
    from src.electra.torch_infer import TorchInfer
    sample = TorchInfer(model_dir)
    vec1, _ = sample.infer(contents)

    ## ONNX infer
    from src.electra.onnx_infer import ONNXInfer
    sample = ONNXInfer(model_dir)
    vec2, _ = sample.infer(contents)

    ## check result data
    values = np.linalg.norm(vec1[0]) * np.linalg.norm(vec2[0])
    cos_sim = vec1[0].dot(vec2[0]) / values
    logging.info(f"cosine similarity:{cos_sim}")
    abs_erro = abs(cos_sim - 1.0)
    if abs_erro < 1e-5:
        logging.info(f"electra check result ok")

if __name__ == "__main__":
    Test_T5()
    print()
    Test_electra()